
print('开始导包。。。')
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
print('导包OK！！！')


app = FastAPI()
print('开始加载模型。。。')
model = SentenceTransformer('/gemini/pretrain/bge-large-zh-v1.5')
print('模型加载OK!!!')

class Sentence(BaseModel):
    text: str


@app.get('/') # 接口
def home():
    return 'hello world'


@app.post('/embedding') # 接口
def get_embedding(sentence: Sentence):
    embedding = model.encode(sentence.text, normalize_embeddings=True).tolist()
    return {"text": sentence.text, "embedding": embedding}


if __name__ == '__main__':
    #
    uvicorn.run(app, host='0.0.0.0', port=50072) # 端口